# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Flower command line interface `new` command."""


import io
import re
import zipfile
from enum import Enum
from pathlib import Path
from string import Template
from typing import Annotated

import requests
import typer

from flwr.supercore.constant import PLATFORM_API_URL

from ..utils import (
    is_valid_project_name,
    parse_app_spec,
    prompt_options,
    prompt_text,
    request_download_link,
    sanitize_project_name,
)


class MlFramework(str, Enum):
    """Available frameworks."""

    PYTORCH = "PyTorch"
    TENSORFLOW = "TensorFlow"
    SKLEARN = "sklearn"
    HUGGINGFACE = "HuggingFace"
    JAX = "JAX"
    MLX = "MLX"
    NUMPY = "NumPy"
    XGBOOST = "XGBoost"
    FLOWERTUNE = "FlowerTune"
    BASELINE = "Flower Baseline"
    PYTORCH_LEGACY_API = "PyTorch (Legacy API, deprecated)"


class LlmChallengeName(str, Enum):
    """Available LLM challenges."""

    GENERALNLP = "GeneralNLP"
    FINANCE = "Finance"
    MEDICAL = "Medical"
    CODE = "Code"


class TemplateNotFound(Exception):
    """Raised when template does not exist."""


def load_template(name: str) -> str:
    """Load template from template directory and return as text."""
    tpl_dir = (Path(__file__).parent / "templates").absolute()
    tpl_file_path = tpl_dir / name

    if not tpl_file_path.is_file():
        raise TemplateNotFound(f"Template '{name}' not found")

    with open(tpl_file_path, encoding="utf-8") as tpl_file:
        return tpl_file.read()


def render_template(template: str, data: dict[str, str]) -> str:
    """Render template."""
    tpl_file = load_template(template)
    tpl = Template(tpl_file)
    if ".gitignore" not in template:
        return tpl.substitute(data)
    return tpl.template


def create_file(file_path: Path, content: str) -> None:
    """Create file including all nessecary directories and write content into file."""
    file_path.parent.mkdir(exist_ok=True)
    file_path.write_text(content, encoding="utf-8")


def render_and_create(file_path: Path, template: str, context: dict[str, str]) -> None:
    """Render template and write to file."""
    content = render_template(template, context)
    create_file(file_path, content)


def print_success_prompt(
    package_name: str, llm_challenge_str: str | None = None
) -> None:
    """Print styled setup instructions for running a new Flower App after creation."""
    prompt = typer.style(
        "🎊 Flower App creation successful.\n\n"
        "To run your Flower App, first install its dependencies:\n\n",
        fg=typer.colors.GREEN,
        bold=True,
    )

    _add = "	huggingface-cli login\n" if llm_challenge_str else ""

    prompt += typer.style(
        f"	cd {package_name} && pip install -e .\n" + _add + "\n",
        fg=typer.colors.BRIGHT_CYAN,
        bold=True,
    )

    prompt += typer.style(
        "then, run the app:\n\n ",
        fg=typer.colors.GREEN,
        bold=True,
    )

    prompt += typer.style(
        "\tflwr run .\n\n",
        fg=typer.colors.BRIGHT_CYAN,
        bold=True,
    )

    prompt += typer.style(
        "💡 Check the README in your app directory to learn how to\n"
        "customize it and how to run it using the Deployment Runtime.\n",
        fg=typer.colors.GREEN,
        bold=True,
    )

    print(prompt)


# Security: prevent zip-slip
def _safe_extract_zip(zf: zipfile.ZipFile, dest_dir: Path) -> None:
    """Extract ZIP file into destination directory."""
    dest_dir = dest_dir.resolve()

    def _is_within_directory(base: Path, target: Path) -> bool:
        try:
            target.relative_to(base)
            return True
        except ValueError:
            return False

    for member in zf.infolist():
        # Skip directory placeholders;
        # ZipInfo can represent them as names ending with '/'.
        if member.is_dir():
            target_path = (dest_dir / member.filename).resolve()
            if not _is_within_directory(dest_dir, target_path):
                raise ValueError(f"Unsafe path in zip: {member.filename}")
            target_path.mkdir(parents=True, exist_ok=True)
            continue

        # Files
        target_path = (dest_dir / member.filename).resolve()
        if not _is_within_directory(dest_dir, target_path):
            raise ValueError(f"Unsafe path in zip: {member.filename}")

        # Ensure parent exists
        target_path.parent.mkdir(parents=True, exist_ok=True)

        # Extract
        with zf.open(member, "r") as src, open(target_path, "wb") as dst:
            dst.write(src.read())


def _download_zip_to_memory(presigned_url: str) -> io.BytesIO:
    """Download ZIP file from Platform API to memory."""
    try:
        r = requests.get(presigned_url, timeout=60)
        r.raise_for_status()
    except requests.RequestException as e:
        typer.secho(
            f"ZIP download failed: {e}",
            fg=typer.colors.RED,
            err=True,
        )
        raise typer.Exit(code=1) from e

    buf = io.BytesIO(r.content)
    # Validate it's a zip
    if not zipfile.is_zipfile(buf):
        typer.secho(
            "Downloaded file is not a valid ZIP",
            fg=typer.colors.RED,
            err=True,
        )
        raise typer.Exit(code=1)
    buf.seek(0)
    return buf


def download_remote_app_via_api(app_spec: str) -> None:
    """Download App from Platform API."""
    # Validate app version and ID format
    app_id, app_version = parse_app_spec(app_spec)
    app_name = app_id.split("/")[1]

    project_dir = Path.cwd() / app_name
    if project_dir.exists():
        if not typer.confirm(
            typer.style(
                f"\n💬 {app_name} already exists, do you want to override it?",
                fg=typer.colors.MAGENTA,
                bold=True,
            )
        ):
            return

    print(
        typer.style(
            f"\n🔗 Requesting download link for {app_id}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    # Fetch ZIP downloading URL
    url = f"{PLATFORM_API_URL}/hub/fetch-zip"
    presigned_url = request_download_link(app_id, app_version, url, "zip_url")

    print(
        typer.style(
            "⬇️  Downloading ZIP into memory...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    zip_buf = _download_zip_to_memory(presigned_url)

    print(
        typer.style(
            f"📦 Unpacking into {project_dir}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    with zipfile.ZipFile(zip_buf) as zf:
        _safe_extract_zip(zf, Path.cwd())

    print_success_prompt(app_name)


# pylint: disable=too-many-locals,too-many-branches,too-many-statements
def new(
    app_name: Annotated[
        str | None,
        typer.Argument(
            help="Flower app name. For remote apps, use the format "
            "'@account_name/app_name' or '@account_name/app_name==x.y.z'. "
            "Version is optional (defaults to latest)."
        ),
    ] = None,
    framework: Annotated[
        MlFramework | None,
        typer.Option(case_sensitive=False, help="The ML framework to use"),
    ] = None,
    username: Annotated[
        str | None,
        typer.Option(case_sensitive=False, help="The Flower username of the author"),
    ] = None,
) -> None:
    """Create new Flower App."""
    if app_name is None:
        app_name = prompt_text("Please provide the app name")

    # Download remote app
    if app_name and app_name.startswith("@"):
        download_remote_app_via_api(app_name)
        return

    if not is_valid_project_name(app_name):
        app_name = prompt_text(
            "Please provide a name that only contains "
            "characters in {'-', a-zA-Z', '0-9'}",
            predicate=is_valid_project_name,
            default=sanitize_project_name(app_name),
        )

    # Set project directory path
    package_name = re.sub(r"[-_.]+", "-", app_name).lower()
    import_name = package_name.replace("-", "_")
    project_dir = Path.cwd() / package_name

    if project_dir.exists():
        if not typer.confirm(
            typer.style(
                f"\n💬 {app_name} already exists, do you want to override it?",
                fg=typer.colors.MAGENTA,
                bold=True,
            )
        ):
            return

    if username is None:
        username = prompt_text("Please provide your Flower username")

    if framework is not None:
        framework_str = str(framework.value)
    else:
        framework_str = prompt_options(
            "Please select ML framework by typing in the number",
            [mlf.value for mlf in MlFramework],
        )

    llm_challenge_str = None
    if framework_str == MlFramework.FLOWERTUNE:
        llm_challenge_value = prompt_options(
            "Please select LLM challenge by typing in the number",
            sorted([challenge.value for challenge in LlmChallengeName]),
        )
        llm_challenge_str = llm_challenge_value.lower()

    if framework_str == MlFramework.BASELINE:
        framework_str = "baseline"

    if framework_str == MlFramework.PYTORCH_LEGACY_API:
        framework_str = "pytorch_legacy_api"

    print(
        typer.style(
            f"\n🔨 Creating Flower App {app_name}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )

    context = {
        "framework_str": framework_str,
        "import_name": import_name.replace("-", "_"),
        "package_name": package_name,
        "project_name": app_name,
        "username": username,
    }

    template_name = framework_str.lower()

    # List of files to render
    if llm_challenge_str:
        files = {
            ".gitignore": {"template": "app/.gitignore.tpl"},
            "pyproject.toml": {"template": f"app/pyproject.{template_name}.toml.tpl"},
            "README.md": {"template": f"app/README.{template_name}.md.tpl"},
            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
            f"{import_name}/server_app.py": {
                "template": "app/code/flwr_tune/server_app.py.tpl"
            },
            f"{import_name}/client_app.py": {
                "template": "app/code/flwr_tune/client_app.py.tpl"
            },
            f"{import_name}/models.py": {
                "template": "app/code/flwr_tune/models.py.tpl"
            },
            f"{import_name}/dataset.py": {
                "template": "app/code/flwr_tune/dataset.py.tpl"
            },
            f"{import_name}/strategy.py": {
                "template": "app/code/flwr_tune/strategy.py.tpl"
            },
        }

        # Challenge specific context
        fraction_train = "0.2" if llm_challenge_str == "code" else "0.1"
        if llm_challenge_str == "generalnlp":
            challenge_name = "General NLP"
            num_clients = "20"
            dataset_name = "flwrlabs/alpaca-gpt4"
        elif llm_challenge_str == "finance":
            challenge_name = "Finance"
            num_clients = "50"
            dataset_name = "flwrlabs/fingpt-sentiment-train"
        elif llm_challenge_str == "medical":
            challenge_name = "Medical"
            num_clients = "20"
            dataset_name = "flwrlabs/medical-meadow-medical-flashcards"
        else:
            challenge_name = "Code"
            num_clients = "10"
            dataset_name = "flwrlabs/code-alpaca-20k"

        context["llm_challenge_str"] = llm_challenge_str
        context["fraction_train"] = fraction_train
        context["challenge_name"] = challenge_name
        context["num_clients"] = num_clients
        context["dataset_name"] = dataset_name
    else:
        files = {
            ".gitignore": {"template": "app/.gitignore.tpl"},
            "README.md": {"template": "app/README.md.tpl"},
            "pyproject.toml": {"template": f"app/pyproject.{template_name}.toml.tpl"},
            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
            f"{import_name}/server_app.py": {
                "template": f"app/code/server.{template_name}.py.tpl"
            },
            f"{import_name}/client_app.py": {
                "template": f"app/code/client.{template_name}.py.tpl"
            },
        }

        # Depending on the framework, generate task.py file
        frameworks_with_tasks = [
            MlFramework.PYTORCH.value,
            MlFramework.JAX.value,
            MlFramework.HUGGINGFACE.value,
            MlFramework.MLX.value,
            MlFramework.TENSORFLOW.value,
            MlFramework.SKLEARN.value,
            MlFramework.NUMPY.value,
            MlFramework.XGBOOST.value,
            "pytorch_legacy_api",
        ]
        if framework_str in frameworks_with_tasks:
            files[f"{import_name}/task.py"] = {
                "template": f"app/code/task.{template_name}.py.tpl"
            }

        if framework_str == "pytorch_legacy_api":
            # Use custom __init__ that better captures name of framework
            files[f"{import_name}/__init__.py"] = {
                "template": f"app/code/__init__.{framework_str}.py.tpl"
            }

        if framework_str == "baseline":
            # Include additional files for baseline template
            for file_name in ["model", "dataset", "strategy", "utils", "__init__"]:
                files[f"{import_name}/{file_name}.py"] = {
                    "template": f"app/code/{file_name}.{template_name}.py.tpl"
                }

            # Replace README.md
            files["README.md"]["template"] = f"app/README.{template_name}.md.tpl"

            # Add LICENSE
            files["LICENSE"] = {"template": "app/LICENSE.tpl"}

    for file_path, value in files.items():
        render_and_create(
            file_path=project_dir / file_path,
            template=value["template"],
            context=context,
        )

    print_success_prompt(package_name, llm_challenge_str)
